数据的定义

from torch.utils.data import  Dataset
from PIL import Image
import  os
class Mydata(Dataset):
    def __init__(self,root_dir,label_dir):
        self.root_dir=root_dir
        self.label_dir=label_dir
        self.path=os.path.join(self.root_dir,self.label_dir)
        self.img_path_list=os.listdir(self.path)

    def __getitem__(self, idx):
        img_name=self.img_path_list[idx]
        img_path=os.path.join(self.path,img_name)
        img=Image.open(img_path)
        label=self.label_dir
        return img,label
    def __len__(self):
        return len(self.img_path_list)
root_dir="D:\\AIlearn\\dataset\\train"
label_dir="ants"
Ants_data=Mydata(root_dir,label_dir)
img,label=Ants_data[0]
# img.show()
复制